#include "Pin.h"
#include "Mat6x6.h"
#include "RigidBodyNode.h"
#include "PrescribedMotion.h"
#include "ABAData.h"
#include "EOMData.h"
#include "State.h"

using namespace RSIM;

///////////////////////////////////////////////////////
// Pin Joint
///////////////////////////////////////////////////////

int Pin::getJointDOF() const{return 1;}

///////////////////////////////////////////////////////

void Pin::updRelVelocity(const State& state, Vect3& w, Vect3& v) const{
        w(2) = state.getUForNode(this->NodeIndex_)[0];
}

///////////////////////////////////////////////////////

void Pin::calc_H_x_UDot(const State& state, double *h_x_udot) const{
        h_x_udot[2] = state.getUDotForNode(this->NodeIndex_)[0];
}

///////////////////////////////////////////////////////

void Pin::setDefaultQ(double *q)const{
        q[0] = 0.0;
}

///////////////////////////////////////////////////////

void Pin::setDefaultU(double *u)const{
        u[0] = 0.0;
}

///////////////////////////////////////////////////////

void Pin::calcQDot(State& state)const{
        state.wgetQDotForNode(this->NodeIndex_)[0] = state.wgetUForNode(this->NodeIndex_)[0];
}

///////////////////////////////////////////////////////

Pin::Pin(const int& NodeIndex):Joint(PinJoint){
        this->NodeIndex_ = NodeIndex;
}

///////////////////////////////////////////////////////

Pin* Pin::New(const int& NodeIndex){
        return new Pin(NodeIndex);
}

///////////////////////////////////////////////////////

void Pin::printJointType() const{
        cout<<"Joint Type=Pin\n";
}

///////////////////////////////////////////////////////

void Pin::updJointTransform(const double *q,
                                const RigidBodyNode *Parent,
                                const int& pChildID,
                                double *betak){
        this->T_MF_.wgetR().setRotationZAxis(q[0]);	
}

///////////////////////////////////////////////////////

void Pin::updVelocity(const State& state, 
                        const RigidBodyNode *Parent, 
                        const Rotation& prF_C_F, 
                        const double *betak,
                        Vect3& w, 
                        Vect3& v) const {
        
        const Vect3& pr_w = Parent->data_->getAngularVelocity();
        const Vect3& pr_v = Parent->data_->getLinearVelocity();
        
        // 1. F_C_prF*pr_w is the absolute angular velocity
        // of the parent expressed in the joint's F frame
        // 2. Vect3(0,0,U_[0]) is the angular velocity of this body
        // which is already expressed in the joint's F frame
        
        double *wptr = w.wPtr();
        FastOps::transpose_mult_mat33_vec3(prF_C_F.Ptr(),pr_w.Ptr(),wptr);
        wptr[2] += state.getUForNode(this->NodeIndex_)[0];
        
        double v_tmp[3];
        
        FastOps::cross(pr_w.Ptr(),betak,v_tmp);
        v_tmp[0] += pr_v(0);
        v_tmp[1] += pr_v(1);
        v_tmp[2] += pr_v(2);
        
        FastOps::transpose_mult_mat33_vec3(prF_C_F.Ptr(), v_tmp, v.wPtr());
}

///////////////////////////////////////////////////////

void Pin::calcAKT(	const RigidBodyNode *Parent, 
                        const RigidBodyNode *Body,
                        const int& ChildID,     
                        const double *betak,
                        State& state,     
                        Vect6& sAkt) const{
        this->calcQDot(state);
        double *UDot = state.wgetUDotForNode(this->NodeIndex_);

        if(Parent){		
                const double *U = state.getUForNode(this->NodeIndex_);
                
                // AKT is calculated in F basis of this body
                double *akt_ptr = sAkt.wPtr();
                const Rotation& prFCF = Body->data_->getRotationFCF();
                
                // Angular velocity of parent expressed in F of parent
                const double *parent_ang_ptr = Parent->data_->getAngularVelocity().Ptr();
                
                {
                        double wwr[3];
                        // cross( w^{k-1}, cross(w^{k-1}, \beta^{k-1}) ) {expressed in F of parent}
                        FastOps::wxwxr(parent_ang_ptr, betak, wwr);
                        
                        // cross( w^{k-1} , w^k_k u^k ) {expressed in F of this body}
                        FastOps::transpose_mult_mat33_vec3(prFCF.Ptr(), wwr, akt_ptr+3);						
                }
                {
                        double ang_vel[3];
                        // Covert Angular velocity of parent to F basis of this body
                        FastOps::transpose_mult_mat33_vec3(prFCF.Ptr(), parent_ang_ptr, ang_vel);
                        
                        akt_ptr[0] = U[0]*ang_vel[1];
                        akt_ptr[1] = -U[0]*ang_vel[0];
                        akt_ptr[2] = 0.0;
                }
                
                // In this case, Akt = Akt + p^k_k \dot{u}^k
                if(this->isMotionPrescribedAndActive_){
                        // Update QDot_ cache from prescribed motion
                        this->PM_->setUDot(state.getSysData(),UDot);						
                        akt_ptr[2] += UDot[0];
                }
        }
        else{
                if(this->isMotionPrescribedAndActive_){
                        this->PM_->setUDot(state.getSysData(), UDot);
                }
        }
}


///////////////////////////////////////////////////////

void Pin::sweepBackwards(	const bool& isLeaf, 
                                const int& ChildID,
                                const Vect6& sAkt,
                                const Rotation& FCF,					    
                                const EOMData *eom,
                                const double *betak,
                                RigidBodyNode *Parent,
                                ABAData *aba) const{
        if(Parent){
                Mat6x6 *Ik3 = aba->wgetIk3();
                Vect6 *Fk3 = aba->wgetFk3();	
                double sIk3[36], sFk3[6];
                
                // cFk = Fk3 - Ik3*A^{k-1}_t
                this->calc_cFk((Fk3)->Ptr(),Ik3->Ptr(),sAkt.Ptr(),aba->wgetcFk()->wPtr());
                                
                if(!isMotionPrescribedAndActive_){		
                        double triang_ik3[36],triang_cfk[6];

                        this->calc_Triang_x_Ik3_RotZ(Ik3->Ptr(),triang_ik3);
                        this->calc_Triang_x_cFk_RotZ(Ik3->Ptr(),aba->getcFk()->Ptr(),triang_cfk);

                        FastOps::spatial_basis_shift_mat(FCF.Ptr(), triang_ik3, sIk3);
                        FastOps::spatial_basis_shift_vec(FCF.Ptr(), triang_cfk, sFk3);
                }
                else{
                        FastOps::spatial_basis_shift_mat(FCF.Ptr(), Ik3->Ptr(), sIk3);
                        FastOps::spatial_basis_shift_vec(FCF.Ptr(), aba->getcFk()->Ptr(), sFk3);	
                }

                double Ik3Child[36], Fk3Child[6];
                this->calc_skbeta_x_Ik3_x_tskbeta(sIk3, betak, Ik3Child);
                this->calc_skbeta_x_cFk(sFk3, betak, Fk3Child);

                Parent->data_->wgetABAData()->wgetIk3()->add(Ik3Child);
                Parent->data_->wgetABAData()->wgetFk3()->add(Fk3Child);
        }// If (Parent)
}

///////////////////////////////////////////////////////

void Pin::sweepForward(const RigidBodyNode *Parent,
                        const int& ChildID,
                        const Vect6& sAkt,
                        const Rotation& FCF,
                        const ABAData *aba,
                        const double *betak,
                        EOMData *eom,
                        State& state
                        ){
        double *UDot = state.wgetUDotForNode(this->NodeIndex_);
        const double *ik3_ptr = aba->getIk3()->Ptr();
        const double *fk3_ptr = aba->getFk3()->Ptr();
        const double *cfk_ptr = aba->getcFk()->Ptr();

        Vect6 *Ak1 = eom->wgetAk1();
        double *ak1_ptr = Ak1->wPtr();
        
        if(Parent){
                const double *Parent_Ak1_Ptr = Parent->data_->getEOMData()->getAk1()->Ptr();
                {
                        double ak1_tmp[6];
                        FastOps::transpose_skbeta_x_V(Parent_Ak1_Ptr, betak, ak1_tmp);
                        FastOps::transpose_spatial_basis_shift_vec(FCF.Ptr(),ak1_tmp, ak1_ptr);
                }

                if(!isMotionPrescribedAndActive_){

                        UDot[0] =  ak1_ptr[0]*ik3_ptr[12]-cfk_ptr[2]+ ak1_ptr[1]*ik3_ptr[13] + ak1_ptr[2]*ik3_ptr[14]; 
                        UDot[0] += ak1_ptr[3]*ik3_ptr[15]+ak1_ptr[4]*ik3_ptr[16]+ak1_ptr[5]*ik3_ptr[17];
                        UDot[0] /= -ik3_ptr[14];

                        ak1_ptr[2] += UDot[0];
                }
                Ak1->add(sAkt.Ptr());		
        } // if(Parent)
        else{
                if(!isMotionPrescribedAndActive_){
                        // calc Udot			
                        UDot[0] = fk3_ptr[2]/ik3_ptr[14];
                        
                        // If motion is prescribed, Udot is already set in call to Pin::calcAKT
                }
                // Spatial acceleration of the F frame of root body.
                double ak1_tmp[]={0.0,0.0,UDot[0],0.0,0.0,0.0};
                Ak1->copy(ak1_tmp);		
        }
        this->calc_m_cFk(fk3_ptr, ik3_ptr, ak1_ptr, eom->wgetFkc1()->wPtr());	
}

///////////////////////////////////////////////////////